// 使用比较简单的方式完成
// 在 leetcode 中会直接超时
var kSmallestPairs = function (nums1, nums2, k) {
  // 生成所有可能的数对
  const pairs = [];
  for (let i = 0; i < nums1.length; i++) {
    for (let j = 0; j < nums2.length; j++) {
      pairs.push([nums1[i], nums2[j]]);
    }
  }

  // 按数对的和排序
  pairs.sort((a, b) => (a[0] + a[1]) - (b[0] + b[1]));

  // 返回前k个元素，如果不足k个则返回全部
  return pairs.slice(0, k);
};

var kSmallestPairs = function (nums1, nums2, k) {
  if (nums1.length === 0 || nums2.length === 0 || k === 0) return [];

  const minHeap = new MinHeap();
  const visited = new Set();

  // 初始化堆，加入第一个元素
  minHeap.insert({ sum: nums1[0] + nums2[0], i: 0, j: 0 });
  visited.add('0,0');

  const result = [];

  while (minHeap.size() > 0 && result.length < k) {
    const { i, j } = minHeap.extractMin();
    result.push([nums1[i], nums2[j]]);

    // 扩展候选对 (i+1, j)
    if (i + 1 < nums1.length && !visited.has(`${i + 1},${j}`)) {
      minHeap.insert({ sum: nums1[i + 1] + nums2[j], i: i + 1, j });
      visited.add(`${i + 1},${j}`);
    }

    // 扩展候选对 (i, j+1)
    if (j + 1 < nums2.length && !visited.has(`${i},${j + 1}`)) {
      minHeap.insert({ sum: nums1[i] + nums2[j + 1], i, j: j + 1 });
      visited.add(`${i},${j + 1}`);
    }
  }

  return result;
};

// 自定义最小堆实现
class MinHeap {
  constructor() {
    this.heap = [];
  }

  size() {
    return this.heap.length;
  }

  insert(element) {
    this.heap.push(element);
    this.bubbleUp();
  }

  bubbleUp() {
    let index = this.heap.length - 1;
    while (index > 0) {
      let parentIndex = Math.floor((index - 1) / 2);
      if (this.heap[parentIndex].sum <= this.heap[index].sum) break;
      [this.heap[parentIndex], this.heap[index]] = [this.heap[index], this.heap[parentIndex]];
      index = parentIndex;
    }
  }

  extractMin() {
    const min = this.heap[0];
    const end = this.heap.pop();
    if (this.heap.length > 0) {
      this.heap[0] = end;
      this.sinkDown();
    }
    return min;
  }

  sinkDown() {
    let index = 0;
    const length = this.heap.length;
    const element = this.heap[0];

    while (true) {
      let leftChildIndex = 2 * index + 1;
      let rightChildIndex = 2 * index + 2;
      let leftChild, rightChild;
      let swap = null;

      if (leftChildIndex < length) {
        leftChild = this.heap[leftChildIndex];
        if (leftChild.sum < element.sum) {
          swap = leftChildIndex;
        }
      }

      if (rightChildIndex < length) {
        rightChild = this.heap[rightChildIndex];
        if (
          (swap === null && rightChild.sum < element.sum) ||
          (swap !== null && rightChild.sum < leftChild.sum)
        ) {
          swap = rightChildIndex;
        }
      }

      if (swap === null) break;
      [this.heap[index], this.heap[swap]] = [this.heap[swap], this.heap[index]];
      index = swap;
    }
  }
}

/**
 * @param {number[]} nums1
 * @param {number[]} nums2
 * @param {number} k
 * @return {number[][]}
 */
// 二分查找和双指针的方式解决
// 如果是记忆的话，直接使用这个就行，上面的最小堆的形式不用记了
var kSmallestPairs = function (nums1, nums2, k) {
  const m = nums1.length
  const n = nums2.length
  let left = nums1[0] + nums2[0]
  let right = nums1[m - 1] + nums2[n - 1]
  let pairSum = right
  while (left <= right) {
    const mid = left + ((right - left) >> 1)
    let cnt = 0;
    let start = 0;
    let end = n - 1;

    while (start < m && end >= 0) {
      if ((nums1[start] + nums2[end]) > mid) {
        end--
      } else {
        cnt += end + 1
        start++
      }
    }
    if (cnt < k) {
      left = mid + 1
    } else {
      pairSum = mid;
      right = mid - 1
    }
  }
  const ans = []
  let pos = n - 1
  // 找到小于目标值pairSum的数对
  for (let i = 0; i < m; i++) {
    while (pos >= 0 && (nums1[i] + nums2[pos]) >= pairSum) {
      pos--
    }
    for (let j = 0; j <= pos && k > 0; j++, k--) {
      const list = [nums1[i], nums2[j]]
      ans.push(list)
    }
  }
  //找到等于目标值pairSum的数对
  pos = n - 1
  for (let i = 0; i < m && k > 0; i++) {
    let start1 = i
    while (i < m - 1 && nums1[i] === nums1[i + 1]) {
      i++
    }
    while (pos >= 0 && nums1[i] + nums2[pos] > pairSum) {
      pos--
    }
    let start2 = pos
    while (pos > 0 && nums2[pos] == nums2[pos - 1]) {
      pos--
    }
    if (nums1[i] + nums2[pos] !== pairSum) {
      continue
    }
    let count = Math.min(k, (i - start1 + 1) * (start2 - pos + 1))
    for (let j = 0; j < count && k > 0; j++, k--) {
      const list = [nums1[i], nums2[pos]];
      ans.push(list);
    }
  }
  return ans
};